"""
File: 批量处理页眉页脚
Author: LiangxuCHEN
Data: 2024/6/29
"""

from docx import Document
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.shared import Pt
from docx.shared import Inches
import os
from enum import Enum

# 利用Enum创建可修改参数枚举类型
class DocChangeKey(Enum):
    """
    定义一个枚举类，用于表示页眉和页脚的对齐方式。
    """
    # 页眉文字
    HEADER_TEXT = 1
    # 页脚文字
    FOOTER_TEXT = 2
    # 页眉图片
    HEADER_IMAGE = 3
    # 页脚图片
    FOOTER_IMAGE = 4

def update_header_footer(doc, save_file_name, **kwargs):
    """
    更新文档的页眉和页脚。
    该函数用于修改指定文档的页眉和页脚内容。通过传入不同的参数，可以实现对页眉和页脚的自定义更新。
    参数:
    - doc: Word文档对象，表示需要进行更新的文档。
    - save_file_name: 字符串，表示更新后的文档保存路径及文件名。
    - **kwargs: 可变关键字参数，用于传递更新页眉和页脚所需的额外参数。
    返回值:
    无返回值，直接在原文档上进行修改并保存。
    """
    # 获取文档的第一个节（通常文档默认至少有一个节）
    section = doc.sections[0]
    header_paragraph = None
    footer_paragraph = None
    # 这里既可以添加图片也可以添加文字
    if h_text:= kwargs.get(DocChangeKey.HEADER_TEXT.name, None):
        # 获取当前节的页眉对象
        header = section.header
        # 获取页眉中的第一个段落对象，通常页眉页脚的编辑都是在段落级别进行
        header_paragraph = header.paragraphs[0]
        # 这里是添加文字的代码
        header_paragraph.text = h_text
        header_paragraph.style.font.size = Pt(12)
        header_paragraph.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
    # 检查是否提供了头部Logo，如果存在，则将其添加到文档中
    if h_logo := kwargs.get(DocChangeKey.HEADER_IMAGE.name, None):
        if header_paragraph is None:
            header = section.header
            header_paragraph = header.paragraphs[0]
        # 创建一个新的运行（run），用于在段落中添加图片
        run = header_paragraph.add_run()
        # 设置图片路径
        logo_path = h_logo
        # 添加图片到运行中，并设置图片宽度
        run.add_picture(logo_path, width=Inches(1))

    # 检查是否提供了页脚文本，如果存在，则对其进行设置
    if f_text := kwargs.get(DocChangeKey.FOOTER_TEXT.name, None):
        # 获取当前节的页脚
        footer = section.footer
        # 获取页脚的第一个段落
        footer_paragraph = footer.paragraphs[0]
        # 设置页脚文本内容
        footer_paragraph.text = f_text
        # 设置段落字体大小
        footer_paragraph.style.font.size = Pt(12)
        # 设置段落对齐方式为居中
        footer_paragraph.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
    if f_logo := kwargs.get(DocChangeKey.FOOTER_IMAGE.name, None):
        if footer_paragraph is None:
            footer = section.footer
            footer_paragraph = footer.paragraphs[0]
        # 创建一个新的运行（run），用于在段落中添加图片
        run = footer_paragraph.add_run()
        # 设置图片路径
        logo_path = f_logo
        # 添加图片到运行中，并设置图片宽度
        run.add_picture(logo_path, width=Inches(1))
    # 保存文档到指定文件名
    doc.save(save_file_name)



def update_all_doc(in_path, out_path, **kwargs):
    """
    更新指定目录下所有docx文档的头部和尾部信息。

    参数:
    in_path: 字符串，输入文档的目录路径。
    out_path: 字符串，处理后文档的输出目录路径。
    **kwargs: 可变关键字参数，用于传递给update_header_footer函数的额外参数。
    """
    # 创建输入路径，如果不存在
    os.makedirs(in_path, exist_ok=True)
    # 获取当前脚本文件的目录
    path = os.path.dirname(os.path.abspath(__file__))
    # 构建完整的输入路径
    in_path = os.path.join(path, in_path)
    # 遍历输入目录中的所有文件
    for filename in os.listdir(in_path):
        print("待处理：", filename)
        try:
            # 如果文件是以'.docx'结尾
            if filename.endswith('.docx'):
                # 构建完整的文档路径
                doc_path = os.path.join(in_path, filename)
                # 打开文档
                doc = Document(doc_path)
                # 构建输出文件的完整路径
                save_file_name = os.path.join(path, out_path, filename)
                print("输出文件:", save_file_name)
                # 更新文档的头部和尾部信息
                update_header_footer(doc, save_file_name, **kwargs)
        except Exception as e:
            # 如果处理过程中发生异常，打印异常信息并继续处理下一个文件
            print(e)
            continue


if __name__ == "__main__":
    # 定义头部图像文件名
    header_logo = "wxpic.png"
    # 定义文档头部文本
    header_text = 'Mycompany'
    # 定义文档尾部文本
    footer_text = '© 2024 Mycompany. All rights reserved.'
    # 定义输入文档的目录
    input_path = "old_docs"
    # 定义输出更新后文档的目录
    output_path = "new_docs"
    # 调用函数更新所有文档的头部和尾部信息
    update_all_doc(
        input_path, output_path,
        **{
            DocChangeKey.HEADER_TEXT.name:header_text,
            DocChangeKey.HEADER_IMAGE.name:header_logo,
            DocChangeKey.FOOTER_TEXT.name:footer_text,
        }
    )
